#pragma once

#include <iostream>
#include <queue>
#include <vector>
#include <thread>
#include <functional>
#include <mutex>
#include <shared_mutex>
#include <condition_variable>
#include <atomic>
#include <future>


namespace hello {

///////////////////////////////////////
/// @brief 线程池，需指定线程数及等待队列大小
///////////////////////////////////////
class thread_pool
{
public:
    thread_pool() : thread_pool(0, 0) {}

    /// @brief 带参构造
    /// @param _pool_sz 线程数
    /// @param _que_sz 等待队列大小
    thread_pool(int _pool_sz, int _que_sz) :
            _shutdown(false),
            _size(_pool_sz),
            _que_sz(_que_sz)
    {
        for (auto i = 0; i < _size; ++i) {
            _ts.emplace_back([this]() {
                // thread task loop when thread pool executing
                while (true)
                {
                    // wait until task queue is not empty or thread pool was shutdown
                    std::unique_lock<std::mutex> _ul(_que_mut);
                    _cv.wait(_ul, [this]() { return _shutdown || !_task_que.empty(); });

                    if (_shutdown && _task_que.empty())
                    {
                        break;
                    }

                    if (!_task_que.empty())
                    {
                        auto _task = std::move(_task_que.front());
                        _task_que.pop();
                        _ul.unlock();
                        _task();
                    }
                }
            });
        }
    }

    /// @brief 不允许拷贝
    thread_pool(const thread_pool &) = delete;

    /// @brief 不允许移动
    thread_pool(thread_pool &&) = delete;

public:
    ~thread_pool()
    {
        {
            std::lock_guard<std::mutex> lg(_que_mut);
            _shutdown = true;
        }
        _cv.notify_all();
        // 等待所有线程结束
        for (auto &_t: _ts)
        {
            _t.join();
        }
    }

public:
    /// @brief 不允许拷贝
    thread_pool &operator=(const thread_pool &) = delete;

    /// @brief 不允许移动
    thread_pool &operator=(thread_pool &&) = delete;

public:
    /// @brief 提交任务到线程池
    /// @tparam Fn 任务函数模板（提交时无需指定）
    /// @tparam ...Args 任务函数对应的参数模板（提交时无需指定）
    /// @param _f 任务函数
    /// @param ..._args 任务函数参数
    /// @return 提交任务返回对应的`std::future`
    template<typename Fn, typename... Args>
    inline auto submit(Fn &&_f, Args &&... _args) -> std::future<decltype(_f(_args...))>
    {
        using Rt = decltype(_f(_args...));

        auto _task = std::make_shared<std::packaged_task<Rt()>>(
                std::bind(
                        std::forward<Fn>(_f),
                        std::forward<Args>(_args)...
                )
        );
        auto fut = _task->get_future();

        // 这里设置为阻塞式提交，后面可以添加超时返回机制
        std::unique_lock<std::mutex> _ul(_que_mut);
        if (_shutdown || _task_que.size() >= _que_sz)
        {
            _ul.unlock();
            if (_shutdown)
                throw std::runtime_error("submit task after thread pool shutdown");
            else
                throw std::runtime_error("submit task after task queue fulled");
        }
        _task_que.emplace([_task]()
        {
            (*_task)();
        });
        _ul.unlock();
        _cv.notify_one();
        return fut;
    }

private:
    bool _shutdown;
    int _size;
    int _que_sz;
    std::vector<std::thread> _ts;
    std::queue<std::packaged_task<void()>> _task_que;
    std::mutex _que_mut;
    std::condition_variable _cv;
};

}
